import torch
import time

def check_pytorch_version():
    """检查PyTorch版本"""
    print("=== PyTorch版本信息 ===")
    print(f"PyTorch版本: {torch.__version__}")
    return torch.__version__

def check_mps_availability():
    """检查MPS可用性"""
    print("\n=== MPS可用性检查 ===")
    
    # 检查是否有mps属性
    has_mps = hasattr(torch.backends, 'mps')
    if not has_mps:
        print("当前PyTorch版本不支持MPS (需要PyTorch 1.12+)")
        print("请升级PyTorch: pip install torch --upgrade")
        return None
    
    # 检查MPS是否可用
    if torch.backends.mps.is_available():
        print("MPS可用，可以使用GPU加速")
        device = torch.device("mps")
    elif torch.backends.mps.is_built():
        print("MPS已构建但不可用，可能是因为macOS版本过低")
        print("需要macOS 12.3+才能使用MPS")
        device = torch.device("cpu")
    else:
        print("MPS未构建，当前PyTorch版本不支持MPS")
        device = torch.device("cpu")
    
    print(f"使用的设备: {device}")
    return device

def check_cuda_availability():
    """检查CUDA可用性（用于兼容NVIDIA GPU）"""
    print("\n=== CUDA可用性检查 ===")
    cuda_available = torch.cuda.is_available()
    print(f"CUDA是否可用: {cuda_available}")
    if cuda_available:
        print(f"CUDA版本: {torch.version.cuda}")
        print(f"GPU数量: {torch.cuda.device_count()}")
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")
    print(f"使用的设备: {device}")
    return device

def demonstrate_computation(device):
    """演示设备计算"""
    print(f"\n=== 使用 {device} 进行计算演示 ===")
    
    # 创建张量
    x = torch.tensor([1., 2., 3.])
    print(f"原始张量: {x} (设备: {x.device})")
    
    # 将张量移到指定设备
    x_device = x.to(device)
    print(f"设备张量: {x_device} (设备: {x_device.device})")
    
    # 在设备上执行计算
    y_device = x_device * 2
    print(f"计算结果: {y_device} (设备: {y_device.device})")
    
    # 将结果移回CPU
    y_cpu = y_device.to("cpu")
    print(f"移回CPU的结果: {y_cpu} (设备: {y_cpu.device})")

def demonstrate_performance_comparison():
    """演示不同设备的性能差异"""
    print("\n=== 设备性能对比 ===")
    
    # 测试CPU性能
    print("测试CPU性能...")
    size = 2000
    x_cpu = torch.randn(size, size)
    y_cpu = torch.randn(size, size)
    
    start_time = time.time()
    z_cpu = torch.mm(x_cpu, y_cpu)
    cpu_time = time.time() - start_time
    print(f"CPU计算时间: {cpu_time:.4f}秒")
    
    # 测试MPS性能（如果可用）
    if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
        print("测试MPS性能...")
        device = torch.device("mps")
        x_mps = x_cpu.to(device)
        y_mps = y_cpu.to(device)
        
        # 预热
        _ = torch.mm(x_mps, y_mps)
        # 在较新的PyTorch版本中使用torch.mps.synchronize()，在旧版本中跳过
        if hasattr(torch, 'mps') and hasattr(torch.mps, 'synchronize'):
            torch.mps.synchronize()
        elif hasattr(torch.backends, 'mps') and hasattr(torch.backends.mps, '_sleep'):
            # 在某些PyTorch版本中使用这个方法等待MPS完成
            torch.cuda.synchronize() if torch.cuda.is_available() else None
        
        start_time = time.time()
        z_mps = torch.mm(x_mps, y_mps)
        # 同样处理同步
        if hasattr(torch, 'mps') and hasattr(torch.mps, 'synchronize'):
            torch.mps.synchronize()
        elif hasattr(torch.backends, 'mps') and hasattr(torch.backends.mps, '_sleep'):
            torch.cuda.synchronize() if torch.cuda.is_available() else None
            
        mps_time = time.time() - start_time
        print(f"MPS计算时间: {mps_time:.4f}秒")
        print(f"MPS相对CPU加速比: {cpu_time/mps_time:.2f}x")
    
    # 测试CUDA性能（如果可用）
    elif torch.cuda.is_available():
        print("测试CUDA性能...")
        device = torch.device("cuda")
        x_cuda = x_cpu.to(device)
        y_cuda = y_cpu.to(device)
        
        # 预热
        _ = torch.mm(x_cuda, y_cuda)
        torch.cuda.synchronize()
        
        start_time = time.time()
        z_cuda = torch.mm(x_cuda, y_cuda)
        torch.cuda.synchronize()
        cuda_time = time.time() - start_time
        print(f"CUDA计算时间: {cuda_time:.4f}秒")
        print(f"CUDA相对CPU加速比: {cpu_time/cuda_time:.2f}x")

def explain_device_info():
    """解释设备相关信息"""
    print("\n=== 设备支持说明 ===")
    print("PyTorch支持多种计算设备:")
    print("  1. CPU - 中央处理器，所有系统都支持")
    print("  2. CUDA - NVIDIA GPU，需要支持CUDA的NVIDIA显卡")
    print("  3. MPS - Metal Performance Shaders，Apple Silicon (M1/M2)芯片")
    
    print("\nMPS相关信息:")
    print("  - 需要PyTorch 1.12+版本")
    print("  - 需要macOS 12.3+系统")
    print("  - 适用于M1、M2等Apple Silicon芯片")
    
    print("\nCUDA相关信息:")
    print("  - 需要支持CUDA的NVIDIA GPU")
    print("  - 需要安装相应的CUDA驱动和工具包")

def main():
    """主函数"""
    print("PyTorch多设备支持演示")
    print("=" * 50)
    
    # 检查PyTorch版本
    version = check_pytorch_version()
    
    # 检查MPS可用性
    mps_device = None
    if hasattr(torch.backends, 'mps'):
        mps_device = check_mps_availability()
    else:
        print("\n当前PyTorch版本无MPS支持，请升级PyTorch")
    
    # 检查CUDA可用性
    cuda_device = check_cuda_availability()
    
    # 根据可用性选择设备进行演示
    if mps_device and mps_device.type == "mps":
        demonstrate_computation(mps_device)
    elif cuda_device.type == "cuda":
        demonstrate_computation(cuda_device)
    else:
        demonstrate_computation(torch.device("cpu"))
    
    # 性能对比
    demonstrate_performance_comparison()
    
    # 设备信息说明
    explain_device_info()
    
    print("\n" + "=" * 50)
    print("程序执行完成!")

if __name__ == '__main__':
    main()